package org.example;

import us.codecraft.webmagic.Page;
import us.codecraft.webmagic.Site;
import us.codecraft.webmagic.Spider;
import us.codecraft.webmagic.processor.PageProcessor;
import us.codecraft.webmagic.selector.Selectable;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static java.util.regex.Pattern.compile;

public class AreaSpider implements PageProcessor {

    /**
     * 公共路径url
     */
    private static String SPIDER_URL = "http://www.stats.gov.cn/sj/tjbz/tjyqhdmhcxhfdm/2022/";


    private Site site = Site.me().setUserAgent("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/80.0.3987.163 Safari/537.36").setTimeOut(1000 * 60).setCharset("UTF-8").setRetryTimes(10).setSleepTime(100);

    /**
     * 爬取数据的年度标识,用于区分批次
     */
    private String year;

    /**
     * 爬取从省级到包含的那一级数据
     * 例如: contain = "city,county";  即爬取省级到县级数据
     */
    private String contain;

    public AreaSpider(String year, String contain) {
        this.year = year;
        this.contain = contain;
    }

    @Override
    public void process(Page page) {
        page.setCharset("UTF-8");
        String url = page.getUrl().toString();
        String[] replaceUrl = url.replace(SPIDER_URL, "").replace(".html", "").split("/");
        if (url.matches(SPIDER_URL + "index.html")) {
            //抓取省级行政区划
            List<Map<String, Object>> province = page.getHtml().xpath("//tr[@class='provincetr']/td").nodes().stream()
                    .filter(selectable -> selectable.xpath("//a/text()") != null)
                    .filter(selectable -> selectable.links().all().size() > 0)
                    .map(selectable -> {
                        String name = selectable.xpath("//a/text()").toString();
                        String newUrl = selectable.links().all().get(0);
                        if (contain.contains("city")) {
                            page.addTargetRequest(newUrl);
                        }
                        String replace = newUrl.replace(SPIDER_URL, "").replace(".html", "");
                        String areaCode = replace + "0000"; //北京市：110000
                        HashMap<String, Object> map = new HashMap<>();
                        map.put("C_NAME", name);
                        map.put("C_CODE", areaCode);
                        map.put("C_LEVEL", "province");
                        map.put("C_CASCADE", "/");//级联
                        map.put("C_FULL_NAME", name);
                        map.put("C_PARENT_CODE", 0L);
                        map.put("C_YEAR", year);
                        return map;
                    }).collect(Collectors.toList());

            page.putField("area", province);
        }

        //抓取市级行政单位
        if (replaceUrl.length == 1 && !replaceUrl[0].equals("index") && contain.contains("city")) {
            List<Map<String, Object>> city = new ArrayList<Map<String, Object>>();
            List<Selectable> cityNodes = page.getHtml().xpath("//tr[@class='citytr']/td").nodes();
            cityNodes.stream().forEach(node -> {
                String name = node.xpath("//a/text()").toString();
                if (!compile("[0-9]*").matcher(name).matches()) {
                    String newUrl = node.links().all().get(0);
                    if (contain.contains("county")) {
                        page.addTargetRequest(newUrl);
                    }
                    String replace = newUrl.replace(SPIDER_URL, "").replace(".html", "");
                    String[] split = replace.split("/");
                    String parentId = split[0] + "0000";
                    String areaCode = split[split.length - 1] + "00";
                    HashMap<String, Object> map = new HashMap<>();
                    map.put("C_NAME", name);
                    map.put("C_CODE", areaCode);
                    map.put("C_LEVEL", "city");
                    map.put("C_CASCADE", "/" + parentId + "/" + areaCode);
                    //C_FULL_NAME 为上一级名称+name
                    map.put("C_FULL_NAME",  name);
                    map.put("C_PARENT_CODE", Long.valueOf(parentId));
                    map.put("C_YEAR", year);
                    city.add(map);
                }
            });
            page.putField("area", city);
        }

        //抓取县级行政单位
        if (replaceUrl.length == 2 && contain.contains("county")) {
            List<Map<String, Object>> county = new ArrayList<Map<String, Object>>();
            List<Selectable> countyNodes = page.getHtml().xpath("//tr[@class='countytr']/td").nodes();
            for (int i = 0; i < countyNodes.size(); i += 2) {
                List<String> code = countyNodes.get(i).xpath("//*/text()").all();
                List<String> name = countyNodes.get(i + 1).xpath("//*/text()").all();
                String areaCode = code.get(0);
                String areaName = name.get(0);
                if (code.size() > 1) {
                    areaCode = code.get(1);
                    areaName = name.get(1);
                    String newUrl = countyNodes.get(i).links().all().get(0);
                    if (contain.contains("town")) {
                        page.addTargetRequest(newUrl);
                    }
                }
                areaCode = areaCode.substring(0, 6);
                String parentId = areaCode.substring(0, 4) + "00";
                HashMap<String, Object> map = new HashMap<>();
                map.put("C_NAME", areaName);
                map.put("C_CODE", areaCode);
                map.put("C_LEVEL", "county");
                map.put("C_CASCADE", "/" + areaCode.substring(0, 2) + "0000/" + parentId + "/" + areaCode);
                map.put("C_FULL_NAME", areaName);
                map.put("C_PARENT_CODE", Long.valueOf(parentId));
                map.put("C_YEAR", year);
                county.add(map);
            }
            page.putField("area", county);
        }

        //抓取镇级行政单位
        if (replaceUrl.length == 3 && contain.contains("town")) {
            List<Map<String, Object>> town = new ArrayList<>();
            List<Selectable> countyNodes = page.getHtml().xpath("//tr[@class='towntr']/td").nodes();
            for (int i = 0; i < countyNodes.size(); i += 2) {
                List<String> code = countyNodes.get(i).xpath("//*/text()").all();
                List<String> name = countyNodes.get(i + 1).xpath("//*/text()").all();
                String areaCode = code.get(0);
                String areaName = name.get(0);
                if (code.size() > 1) {
                    areaCode = code.get(1);
                    areaName = name.get(1);
                    String newUrl = countyNodes.get(i).links().all().get(0);
                    if (contain.contains("village")) {
                        page.addTargetRequest(newUrl);
                    }
                }
                areaCode = areaCode.substring(0, 9);
                String parentId = areaCode.substring(0, 6);
                HashMap<String, Object> map = new HashMap<>();
                map.put("C_NAME", areaName);
                map.put("C_CODE", areaCode);
                map.put("C_LEVEL", "town");
                map.put("C_CASCADE", "/" + areaCode.substring(0, 2) + "0000/" + areaCode.substring(0, 4) + "00/" + parentId + "/" + areaCode);
                map.put("C_FULL_NAME", areaName);
                map.put("C_PARENT_CODE", Long.valueOf(parentId));
                map.put("C_YEAR", year);
                town.add(map);
            }
            page.putField("area", town);
        }

        //抓取乡级行政单位
        if (replaceUrl.length == 4 && contain.contains("village")) {
            List<Map<String, Object>> village = new ArrayList<Map<String, Object>>();
            List<Selectable> countyNodes = page.getHtml().xpath("//tr[@class='villagetr']/td").nodes();
            for (int i = 0; i < countyNodes.size(); i += 3) {
                String areaCode = countyNodes.get(i).xpath("//*/text()").get();
                String areaName = countyNodes.get(i + 2).xpath("//*/text()").get();
                String parentId = areaCode.substring(0, 9);
                HashMap<String, Object> map = new HashMap<>();
                map.put("C_NAME", areaName);
                map.put("C_CODE", areaCode);
                map.put("C_LEVEL", "village");
                map.put("C_CASCADE", "/" + areaCode.substring(0, 2) + "0000/" + areaCode.substring(0, 4) + "00/" + areaCode.substring(0, 6) + "/" + parentId + "/" + areaCode);
                map.put("C_FULL_NAME", areaName);
                map.put("C_PARENT_CODE", Long.valueOf(parentId));
                map.put("C_YEAR", year);
                village.add(map);
            }
            page.putField("area", village);
        }
    }

    @Override
    public Site getSite() {
        return site;
    }

    public static void main(String[] args) {
        Spider.create(new AreaSpider("2022", "city, county")).addUrl("http://www.stats.gov.cn/sj/tjbz/tjyqhdmhcxhfdm/2022/index.html")
                .addPipeline(new SqlPipeline()).thread(16).run();
    }
}
